# -*- coding: utf-8 -*-
"""
extract video frames
"""

import os
import sys
import argparse
import shutil
import cv2


def extract_video_frames(video_full_path, img_save_dir_per_video, sample_freq):
    """
    extract every single video frames
    :param video_full_path:
    :param img_save_dir_per_video:
    :param sample_freq:
    :return:
    """

    if os.path.exists(img_save_dir_per_video):
        shutil.rmtree(img_save_dir_per_video)
        os.makedirs(img_save_dir_per_video)
    else:
        os.mkdir(img_save_dir_per_video)

    video_cap = cv2.VideoCapture(video_full_path)
    if not video_cap.isOpened():
        print("video {} open failed".format(video_full_path))
        return 0

    frame_num = int(video_cap.get(cv2.CAP_PROP_FRAME_COUNT))
    success, frame = video_cap.read()
    while success:
        frame_pos = int(video_cap.get(cv2.CAP_PROP_POS_FRAMES))
        if frame_pos % sample_freq == 0:
            frame_name = "{}_{}.jpg".format(os.path.split(video_full_path)[-1].split(".")[0],
                                            int(frame_pos/sample_freq))
            # frame_name = str(int(frame_pos / sample_freq)) + '.jpg'
            pic_save_path = os.path.join(img_save_dir_per_video, frame_name)
            cv2.imwrite(pic_save_path, frame)
        if int(frame_pos) % 100 == 0:
            print("Processing video: {}, Frame: {}, Total: {} ".format(
                os.path.split(video_full_path)[-1], frame_pos, frame_num))
        success, frame = video_cap.read()
    return frame_num


def batch_extract_video_frames(video_dir, img_save_dir, sample_freq, start_frame_id=0):
    """
    extract all the videos in 'video_dir' and save frame in 'img_save_dir'
    :param video_dir:      video directory
    :param img_save_dir:   image save dir
    :param sample_freq:    save frames every other 'sample_freq'
    :param start_frame_id: save frames from 'start_frame_id'
    :return: None
    """
    files = os.listdir(video_dir)
    videos = []
    for file in files:
        if os.path.splitext(file)[1] in [".mp4", ".avi", ".mkv"]:
            videos.append(file)
    video_num = len(videos)

    for video in videos:
        video_full_path = os.path.join(video_dir, video)
        img_save_dir_per_video = os.path.join(img_save_dir, video.split(".")[0])

        frame_count = extract_video_frames(video_full_path,
                                           img_save_dir_per_video, sample_freq)
        print("Finish video {}, "
              "Total: {}, Extracted: {} frames".format(video, video_num, frame_count))
        print("*"*80)
        print()


if __name__ == '__main__':
    try:
        VIDEO_DIR = sys.argv[1]
        IMG_SAVE_DIR = sys.argv[2]
        SAMPLE_FREQUENCY = sys.argv[3]
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--video_dir", type=str,
                            default=os.path.join(os.getenv("HOME"),
                                                 "PycharmProjects/data_tool/data/video"),
                            help="video directory")
        PARSER.add_argument("--img_save_dir", type=str,
                            default=os.path.join(os.getenv("HOME"),
                                                 "PycharmProjects/data_tool/data/video"),
                            help="image save directory")
        PARSER.add_argument("--sample_frequency", type=int, default=50,
                            help="sampling frequency.")
        ARGS = PARSER.parse_args()
        VIDEO_DIR = ARGS.video_dir
        IMG_SAVE_DIR = ARGS.img_save_dir
        SAMPLE_FREQUENCY = ARGS.sample_frequency
    batch_extract_video_frames(VIDEO_DIR, IMG_SAVE_DIR, SAMPLE_FREQUENCY)
